'''
Created on Aug 28, 2012

@author: Andre Dozier
'''
import numpy as np 
import matplotlib as mlib
import matplotlib.pyplot as plt
import csv
import random
import string
import matplotlib.dates as mdates 
import datetime as dt  


def complement(ind_arr, n):
    """
    Find the complement of the set of indices in ind_arr from
    arange(n)
    """
    mat = np.ones(n)
    np.put(mat, ind_arr, 0)
    out = np.nonzero(mat)
    return out[0]

def plotinputs(basename, Xtrain, Ttrain, XNames, TNames, subplotsize=(),
               indexCol= -1):
    # plot all variables against the index column  
    # individual figures
    cols = range(Xtrain.shape[1])
    if indexCol != -1:
        cols.remove(indexCol) 
    for i in cols:
        newFigure(XNames[indexCol], XNames[i])
        if indexCol == -1:
            plt.plot(Xtrain[:, i], 'k-')
        else:
            plt.plot(Xtrain[:, indexCol], Xtrain[:, i], 'o')
        flname = basename + "_" + XNames[i] + "Tseries.png"
        saveFigure(flname)
    for i in range(Ttrain.shape[1]): 
        newFigure(XNames[indexCol], string.join(TNames, ','))
        if indexCol == -1:
            plt.plot(Ttrain[:, i], 'k-')
        else:
            plt.plot(Xtrain[:, indexCol], Ttrain[:, i], 'o')
        flname = basename + "_" + TNames[i] + "Tseries.png"
        saveFigure(flname)
        
    # one big figure
    if subplotsize != (): 
        fig = plt.figure() 
        mlib.rcParams.update({'font.size': 5})
        for i in cols: 
            ax = fig.add_subplot(subplotsize[0], subplotsize[1],
                                 i + 1)
            h = ax.plot(Xtrain[:, indexCol], Xtrain[:, i], 'o')
            ax.set_xlabel(XNames[indexCol])
            ax.set_ylabel(XNames[i])
            ax.yaxis.set_label_coords(-0.1, 0.5)
            ax.xaxis.set_label_coords(0.5, -0.12)
        flname = basename + "_combinedTseries.png"
        if len(TNames) > 1: 
            plt.figlegend(h, TNames, loc='best')
        saveFigure(flname, newdpi=300) 
        mlib.rcdefaults()  
        
    # plot inputs versus outputs 
    # (works if outputs are similar in range)
    for i in range(Xtrain.shape[1]):
        newFigure(XNames[i], string.join(TNames, ','))
        plt.plot(Xtrain[:, i], Ttrain[:, 0], 'o')
        flname = basename + "_" + XNames[i] + ".png"
        saveFigure(flname) 
            
    # plot all inputs versus outputs in subplots on one figure
    if subplotsize != ():
        fig = plt.figure()
        mlib.rcParams.update({'font.size': 5})
        for i in range(Xtrain.shape[1]):
            ax = fig.add_subplot(subplotsize[0], subplotsize[1], i + 1)
            h = ax.plot(Xtrain[:, i], Ttrain, 'o')
            ax.set_xlabel(XNames[i])
            ax.set_ylabel(string.join(TNames, ','))
            ax.yaxis.set_label_coords(-0.1, 0.5)
            ax.xaxis.set_label_coords(0.5, -0.12)
        flname = basename + "_combined.png"
        if len(TNames) > 1: 
            plt.figlegend(h, TNames, loc='lower right',
                          bbox_to_anchor=(0.9, 0.1))
        saveFigure(flname, newdpi=300) 
        mlib.rcdefaults()
        
    # plot outputs versus each other
    fig = plt.figure()
    mlib.rcParams.update({'font.size': 5})
    nout = len(TNames)
    for i in range(nout): 
        for j in range(nout): 
            ax = fig.add_subplot(nout, nout, i * nout + j + 1)
            if i == j: 
                ax.hist(Ttrain[:, j], 50, normed=1, facecolor='black', alpha=0.75)
            else:
                h = ax.plot(Ttrain[:, j], Ttrain[:, i], 'k.')
                if nout <= 5: 
                    ax.set_ylabel(TNames[i])
            if nout <= 5: 
                ax.set_xlabel(TNames[j])
            ax.yaxis.set_label_coords(-0.1, 0.5)
            ax.xaxis.set_label_coords(0.5, -0.1)
            if nout > 5: 
                ax.set_xticklabels([])
                ax.set_yticklabels([])
            
    flname = basename + "_targets.png"
    saveFigure(flname, newdpi=300)
    mlib.rcdefaults()

def toDateNum(year, month, day):
    return mdates.date2num(dt.datetime(year, month, day))

def toDateNums(years, months, days):
    return np.vectorize(toDateNum)(years, months, days)

def plot_date(years, months, days, Y, Ttrain, filepath="",
              trainI=None, validI=None, testI=None):
    plt.figure()
    datevals = toDateNums(years, months, days) 
    if trainI == None: 
        plt.plot_date(datevals, Ttrain, 'ko', label='Measured')
        plt.plot_date(datevals, Y, 'k-', label='Simulated')
    else: 
        plt.plot_date(datevals, Ttrain, 'k-', label='Measured')
        plt.plot_date(datevals[trainI], Y[trainI], 'x', label='Training')
        if validI != None: 
            plt.plot_date(datevals[validI], Y[validI], 'x', label='Validation')
        if testI != None: 
            plt.plot_date(datevals[testI], Y[testI], 'x', label='Test')
    plt.legend()
    if filepath != "": 
        saveFigure(filepath)

def newFigure(xlabel, ylabel):
    '''
    @summary: 
    Sets up a new figure with given x- and y-axes labels 
    '''
    plt.figure()
    plt.ion()
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    
def saveFigure(outFilePath, newdpi=None):
    '''
    @summary: 
    Saves a figure to the specified output file location (make sure to 
    include a common image format extension in the file name). 
    ''' 
    if outFilePath != "":
        print("Saving figure to " + outFilePath)
        plt.savefig(outFilePath, dpi=newdpi)

def SplitData(Xtrain, Ttrain, TrainingFraction, ValidationFraction=0,
              rndSeed=None, splitCtr=1):
    '''
    @summary: 
    Splits data into training and testing datasets
    
    @note: 
    Use this before Xtrain has a column of ones added to it
    '''
    # set the random seed if desired
    if rndSeed != None: 
        random.seed(rndSeed)
    
    # get indices for training and testing datasets
    nsamples = Xtrain.shape[0]
    allI = xrange(nsamples)
    nTrain = int(round(nsamples * TrainingFraction))
    nValid = int(round(nsamples * ValidationFraction))
    rndRowI = random.sample(allI, nsamples)
    trainI = rndRowI[:nTrain]
    validI = rndRowI[nTrain:nTrain + nValid]
    testI = rndRowI[nTrain + nValid:]

    # Get training datasets 
    XTrain = Xtrain[trainI, :]
    TTrain = Ttrain[trainI, :]
    
    # check for constant columns (linear dependent with bias column) 
    hasZero = (XTrain.std(0) == 0).any()
    if hasZero and splitCtr < 100: 
        splitCtr += 1
        return SplitData(Xtrain, Ttrain,
                         TrainingFraction, ValidationFraction,
                         None, splitCtr)
    hasZero = (TTrain.std(0) == 0).any()
    if hasZero and splitCtr < 100:
        splitCtr += 1  
        return SplitData(Xtrain, Ttrain,
                         TrainingFraction, ValidationFraction,
                         None, splitCtr)
    
    # Get testing and validation datasets 
    XTest = Xtrain[testI, :]
    TTest = Ttrain[testI, :]    
    XValid = Xtrain[validI, :]
    TValid = Ttrain[validI, :]

    return [XTrain, TTrain, XTest,
            TTest, XValid, TValid,
            trainI, testI, validI]

def makeStandardize(vals, axis=0, mean=None, stddev=None):
    '''
    @summary: 
    Standardizes the values and returns functions that standardize and 
    unstandardize any new data
    '''
    if mean == None: 
        mean = vals.mean(axis)
    if stddev == None: 
        stddev = vals.std(axis, ddof=1)
    def standardize(origX):
        return (origX - mean) / stddev 
    def unStandardize(stdX):
        return stddev * stdX + mean
    return (standardize(vals) if vals != None else vals,
            standardize, unStandardize, mean, stddev)

def AddOnes(Xtrain):
    '''
    @summary:
    Adds a column of ones as the first column of the Xtrain matrix if not 
    already present.
    '''
    if not((Xtrain[:, 0:1] == 1).all(0)[0]):
        colOfOnes = np.ones((Xtrain.shape[0], 1))
        return np.hstack((colOfOnes, Xtrain))
    else: 
        return Xtrain

def RemoveOnes(Xtrain):
    '''
    @summary: 
    Removes the column of ones at the first column of the Xtrain matrix if 
    present. 
    '''
    if (Xtrain[:, 0:1] == 1).all(0)[0]: 
        return Xtrain[:, 1:]
    else: 
        return Xtrain

def Clean(vals):
    '''
    @summary:
    Cleans the x-values and y-values from all non-numbers by 
    omitting those data.

    @return: 
    [clean_vals]
    '''
    notNans = (np.isnan(vals) == False)
    goodRowsMask = notNans.all(axis=1)
    return vals[goodRowsMask, :]

def bootstrap(data, func, nboot):
    '''Produce nboot bootstrap samples from applying func to data'''
    n = len(data)
    resamples = np.array([[random.choice(data) for _ in range(n)]
                          for _ in range(nboot)])
    return np.apply_along_axis(func, 1, resamples)

class DataConverter:
    '''
    Converts a specified 'missingValue' into a NaN later to be cleaned 
    up.
    '''
    def __init__(self, missingValue):
        self.missingValue = missingValue
        
    def ConvertMissingToNaN(self, s):
        '''
        @summary: 
        If a single value is a missing value, it returns np.nan. 
        Otherwise, returns the same value converted to a floating point
        number 
        '''
        if s == self.missingValue: 
            return np.nan
        else: 
            return float(s) 

class DataReader:
    '''
    @summary:
    Reads commonly used data text files
    '''
    
    def __init__(self, filepath, skiprows,
                 xcols, ycols, converters=None, delim=","):
        '''
        @summary:
        Constructor
        
        @param filepath: 
        The file in which data is located
        
        @param skiprows: 
        Number of beginning rows to skip in the file
        
        @param xcols: 
        The columns of independent variables
        
        @param ycols: 
        The columns of dependent variables
        
        @param converters: 
        Methods to convert data from within the data file to a floating 
        point number or NaN
        '''
        self.filepath = filepath
        self.skiprows = skiprows
        self.xcols = xcols
        self.ycols = ycols
        self.converters = converters
        self.hasBeenRead = False
        self.delim = delim
  
    def Read(self, PerformClean=True, saveplot=False, subplotsize=()):
        '''
        Reads data from the file and returns an array of the values
        '''
        
        # read headers 
        if self.skiprows != 0:
            for row in csv.reader(open(self.filepath, 'r')): 
                self.XNames = [row[xcol] for xcol in self.xcols]
                self.TNames = [row[ycol] for ycol in self.ycols] 
                break
        
        # read and clean data
        self.allvals = np.loadtxt(fname=self.filepath,
                                  delimiter=self.delim,
                             skiprows=self.skiprows,
                             usecols=self.xcols + self.ycols,
                             converters=self.converters)
        if PerformClean: 
            self.allvals = Clean(self.allvals)

        # split data to x and y (target) values
        ncols = len(self.xcols)
        self.X = self.allvals[:, 0:ncols]
        self.T = self.allvals[:, ncols:]

        if saveplot: 
            base = self.filepath[:len(self.filepath) - 4]
            plotinputs(base, self.X, self.T, self.XNames, self.TNames, subplotsize)
            
        self.hasBeenRead = True
        return self.allvals
         
class Standardizer: 
    '''
    @summary: 
    Houses functions to standardize and unstandardize data
    '''
    def __init__(self, UnstandardizedData=None, means=None, stdevs=None,
                 axis=0):
        if UnstandardizedData != None: 
            (self.StandardizedData,
             self.Stdize,
             self.UnStdize,
             self.Means,
             self.Stdevs) = makeStandardize(UnstandardizedData, axis)
        else: 
            (self.StandardizedData,
             self.Stdize,
             self.UnStdize,
             self.Means,
             self.Stdevs) = makeStandardize(None, axis, means, stdevs)
  
    def Standardize(self, UnstandardizedData):
        return self.Stdize(UnstandardizedData)
    
    def Unstandardize(self, StandardizedData):
        return self.UnStdize(StandardizedData)
